import cv2
import numpy as np
import time
from matplotlib import pyplot as plt

def drawmask(mask,x1,y1,x2,y2,type):
    if(x1 < 0):
        x1 = 0
    if(y1 < 0):
        y1 = 0
    if(x2 > mask.shape[1]):
        x2 = mask.shape[1]
    if(y2 > mask.shape[0]):
        y2 = mask.shape[0]
    mask[y1:y2, x1:x2] = type

img = cv2.imread('../images/IMG_8800.JPG')
# img = cv2.resize(img,(600,800))
img = cv2.resize(img,(int(0.3*img.shape[1]),int(0.3*img.shape[0])))
OLD_IMG = img.copy()
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (3, 3), 0)

SIZE = (1, 65)
bgdModle = np.zeros(SIZE, np.float64)
fgdModle = np.zeros(SIZE, np.float64)

# circles = cv2.HoughCircles(blur, cv2.HOUGH_GRADIENT, 1, 3000,
#                                param1=100, param2=30, minRadius=30, maxRadius=50)

circles = cv2.HoughCircles(blur, cv2.HOUGH_GRADIENT, 1, 3000,
                           param1=100, param2=30,
                           minRadius=int(1/20*min(img.shape[0],img.shape[1])),
                           maxRadius=int(1/12*max(img.shape[0],img.shape[1])))
for i in circles[0, :]:
    # draw the outer circle
    cv2.circle(OLD_IMG, (i[0], i[1]), i[2], (0, 255, 0), 2)
    # draw the center of the circle
    cv2.circle(OLD_IMG, (i[0], i[1]), 2, (0, 0, 255), 3)

coin_cx = circles[0][0][0]
coin_cy = circles[0][0][1]
coin_r = circles[0][0][2]
print(circles)
print(coin_cx,coin_cy,coin_r)

mask = np.zeros(img.shape[:2], np.uint8)
# mask[0:img.shape[0],int(coin_cx-1.4*coin_r):int(coin_cx+1.4*coin_r)] = 2
# mask[int(coin_cy-8*coin_r):int(coin_cy+8*coin_r),
#      int(coin_cx-0.6*coin_r):int(coin_cx+0.6*coin_r)] = 3
# mask[int(coin_cy-2*coin_r):int(coin_cy+2*coin_r),
#      int(coin_cx-0.6*coin_r):int(coin_cx+0.6*coin_r)] = 1
drawmask(mask,int(coin_cx-1.4*coin_r),0,int(coin_cx+1.4*coin_r),mask.shape[0], 2)
drawmask(mask,
         int(coin_cx-0.4*coin_r),int(coin_cy-8*coin_r),
         int(coin_cx+0.4*coin_r),int(coin_cy+8*coin_r), 3)
drawmask(mask,
         int(coin_cx-0.4*coin_r),int(coin_cy-1.4*coin_r),
         int(coin_cx+0.4*coin_r),int(coin_cy+1.4*coin_r), 1)


# ============= For Rectangle Initial Method =============
# rect = (int(circles[0][0][0])-int(1.2*circles[0][0][2]),
#         int(circles[0][0][1])-11*int(circles[0][0][2]),
#         int(2.4*circles[0][0][2]),
#         22*int(circles[0][0][2]))
#
# rect = (int(circles[0][0][0])-int(1.5*circles[0][0][2]),
#         1,
#         int(3*circles[0][0][2]),
#         img.shape[0])

cv2.rectangle(OLD_IMG, (int(circles[0][0][0])-int(1.2*circles[0][0][2]),
                        int(circles[0][0][1])-11*int(circles[0][0][2])),
              (int(circles[0][0][0]) + 1 * int(1.2*circles[0][0][2]),
               int(circles[0][0][1]) + 11 * int(circles[0][0][2])), (0, 255, 0), 2, 4)

time_start=time.time()
cv2.grabCut(img, mask, None, bgdModle, fgdModle, 10, cv2.GC_INIT_WITH_MASK)
time_end = time.time()
print('Totally time: ',time_end-time_start,' (s)')

mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
img *= mask2[:, :, np.newaxis]

plt.subplot(131), plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.title("Grabcut"), plt.xticks([]), plt.yticks([])
plt.subplot(132), plt.imshow(cv2.cvtColor(OLD_IMG, cv2.COLOR_BGR2RGB))
plt.title("Original"), plt.xticks([]), plt.yticks([])
plt.subplot(133), plt.imshow(mask)
plt.title("Mask"), plt.xticks([]), plt.yticks([])

plt.show()